package main

import (
	"github.com/gorilla/websocket"
	"log"
	"net"
	"net/http"
	"strings"
	"sync"
)

var conns = make(map[*websocket.Conn]struct{})
var mutex sync.Mutex

func HandleWebSocket(w http.ResponseWriter, r *http.Request) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("升级成websocket失败...:", err)
		return
	}

	go newConnect(ws, r)
}

func newConnect(ws *websocket.Conn, r *http.Request) {

	// 创建一个 channel 用于发送和接收退出信号
	exitChan := make(chan struct{})
	defer func() {
		close(exitChan)
		defer ws.Close()
	}()

	mutex.Lock()
	conns[ws] = struct{}{}
	mutex.Unlock()

	log.Printf("客户端连接...")

	// 获取远程地址（客户端IP）
	// 尝试从HTTP头部获取真实的客户端IP地址
	clientIP := r.Header.Get("X-Forwarded-For")
	if clientIP == "" {
		clientIP = r.Header.Get("X-Real-IP")
	}
	if clientIP == "" {
		clientIP = r.RemoteAddr // 如果没有代理头部，使用RemoteAddr
	}

	// 获取请求的端口（客户端端口）
	_, clientPort, err := net.SplitHostPort(clientIP)
	if err != nil {
		log.Printf("Error splitting host and port: %v", err)
	}

	log.Printf("新的客户端连接产生，IP: %s, port: %s", strings.TrimLeft(clientIP, "["), clientPort)

	// 处理客户端关闭连接
	ws.SetCloseHandler(func(code int, text string) error {
		mutex.Lock()
		delete(conns, ws) // 从conns中删除当前连接
		mutex.Unlock()

		log.Printf("客户端连接关闭: IP=%s, Code=%d, Text=%s", clientIP, code, text)
		return nil
	})

	//开启协程读取数据并发送数据
	go func() {
		// 循环读取消息
		for {
			select {
			case <-exitChan:
				log.Println("发送循环退出")
				return
			default:
				// 从通道接收数据
				data := <-dataChannel
				if data == nil {
					continue
				}
				// 实时向客户端发送消息
				if err := ws.WriteMessage(websocket.TextMessage, data); err != nil {
					if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
						// 正常关闭或客户端离开
						log.Println("客户端正常关闭连接")
					} else {
						// 其他错误
						log.Println("发送消息时发生错误:", err)
					}
					return
				}
			}

		}
	}()

	// 阻塞循环读取消息（放后面）
	for {
		_, _, err := ws.ReadMessage()
		if err != nil {
			if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
				// 客户端关闭了连接
				log.Println("客户端关闭了连接")
			} else {
				// 处理其他错误
				log.Println("客户端关闭", err)
			}
			// 发送退出信号
			exitChan <- struct{}{}
			break
		}

	}

}
